// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.impala.customcluster;

import static org.apache.impala.testutil.LdapUtil.*;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.net.URLConnection;
import java.util.HashMap;
import java.util.Map;

import com.google.common.collect.Range;
import org.apache.impala.testutil.WebClient;
import org.junit.After;
import org.junit.Test;

/**
 * Tests that Web Server works as expected when JWT authentication is being used.
 */
public class JwtWebserverTest {
  private static final Range<Long> zero = Range.closed(0L, 0L);

  WebClient client_ = new WebClient(TEST_USER_1, TEST_PASSWORD_1);

  public void setUp(String extraArgs, String startArgs) throws Exception {
    Map<String, String> env = new HashMap<>();
    env.put("IMPALA_WEBSERVER_USERNAME", TEST_USER_1);
    env.put("IMPALA_WEBSERVER_PASSWORD", TEST_PASSWORD_1);
    int ret = CustomClusterRunner.StartImpalaCluster(extraArgs, env, startArgs);
    assertEquals(ret, 0);
  }

  @After
  public void cleanUp() throws Exception {
    // Leave a cluster running with the default configuration.
    CustomClusterRunner.StartImpalaCluster();
    client_.close();
  }

  private void verifyAuthMetrics(
      Range<Long> expectedAuthSuccess, Range<Long> expectedAuthFailure, String auth_type)
        throws Exception {
    long actualAuthSuccess =
        (long) client_.getMetric("impala.webserver.total-" + auth_type +
            "-token-auth-success");
    assertTrue("Expected: " + expectedAuthSuccess + ", Actual: " + actualAuthSuccess,
        expectedAuthSuccess.contains(actualAuthSuccess));
    long actualAuthFailure =
        (long) client_.getMetric("impala.webserver.total-" + auth_type +
             "-token-auth-failure");
    assertTrue("Expected: " + expectedAuthFailure + ", Actual: " + actualAuthFailure,
        expectedAuthFailure.contains(actualAuthFailure));
  }

  /**
   * Tests if sessions are authenticated by verifying the JWT token for connections
   * to the Web Server.
   * Since we don't have Java version of JWT library, we use pre-calculated JWT token
   * and JWKS. The token and JWK set used in this test case were generated by using
   * BE unit-test function JwtUtilTest::VerifyJwtRS256.
   */
  @Test
  public void testWebserverJwtAuth() throws Exception {
    String jwksFilename =
        new File(System.getenv("IMPALA_HOME"), "testdata/jwt/jwks_rs256.json").getPath();
    setUp(String.format(
              "--jwt_token_auth=true --jwt_validate_signature=true --jwks_file_path=%s "
                  + "--jwt_allow_without_tls=true",
              jwksFilename),
        "");

    // Case 1: Authenticate with valid JWT Token in HTTP header.
    String jwtToken =
        "eyJhbGciOiJSUzI1NiIsImtpZCI6InB1YmxpYzpjNDI0YjY3Yi1mZTI4LTQ1ZDctYjAxNS1m"
        + "NzlkYTUwYjViMjEiLCJ0eXAiOiJKV1MifQ.eyJpc3MiOiJhdXRoMCIsInVzZXJuYW1lIjoia"
        + "W1wYWxhIn0.OW5H2SClLlsotsCarTHYEbqlbRh43LFwOyo9WubpNTwE7hTuJDsnFoVrvHiWI"
        + "02W69TZNat7DYcC86A_ogLMfNXagHjlMFJaRnvG5Ekag8NRuZNJmHVqfX-qr6x7_8mpOdU55"
        + "4kc200pqbpYLhhuK4Qf7oT7y9mOrtNrUKGDCZ0Q2y_mizlbY6SMg4RWqSz0RQwJbRgXIWSgc"
        + "bZd0GbD_MQQ8x7WRE4nluU-5Fl4N2Wo8T9fNTuxALPiuVeIczO25b5n4fryfKasSgaZfmk0C"
        + "oOJzqbtmQxqiK9QNSJAiH2kaqMwLNgAdgn8fbd-lB1RAEGeyPH8Px8ipqcKsPk0bg";
    attemptConnection("Bearer " + jwtToken, "127.0.0.1");
    verifyAuthMetrics(Range.closed(1L, 1L), zero, "jwt");

    // Case 2: Failed with invalid JWT Token.
    String invalidJwtToken =
        "eyJhbGciOiJSUzI1NiIsImtpZCI6InB1YmxpYzpjNDI0YjY3Yi1mZTI4LTQ1ZDctYjAxNS1m"
        + "NzlkYTUwYjViMjEiLCJ0eXAiOiJKV1MifQ.eyJpc3MiOiJhdXRoMCIsInVzZXJuYW1lIjoia"
        + "W1wYWxhIn0.";
    try {
      attemptConnection("Bearer " + invalidJwtToken, "127.0.0.1");
    } catch (IOException e) {
      assertTrue(e.getMessage().contains("Server returned HTTP response code: 401"));
    }
    verifyAuthMetrics(Range.closed(1L, 1L), Range.closed(1L, 1L), "jwt");

    // Case 3: Failed without "Bearer" token.
    try {
      attemptConnection("Basic VGVzdDFMZGFwOjEyMzQ1", "127.0.0.1");
    } catch (IOException e) {
      assertTrue(e.getMessage().contains("Server returned HTTP response code: 401"));
    }
    // JWT authentication is not invoked.
    verifyAuthMetrics(Range.closed(1L, 1L), Range.closed(1L, 1L), "jwt");

    // Case 4: Failed without "Authorization" header.
    try {
      attemptConnection(null, "127.0.0.1");
    } catch (IOException e) {
      assertTrue(e.getMessage().contains("Server returned HTTP response code: 401"));
    }
    // JWT authentication is not invoked.
    verifyAuthMetrics(Range.closed(1L, 1L), Range.closed(1L, 1L), "jwt");
  }

  /**
   * Tests if sessions are authenticated by verifying the OAuth token for connections
   * to the Web Server.
   * Since we don't have Java version of JWT library, we use pre-calculated JWT token
   * and JWKS. The token and JWK set used in this test case were generated by using
   * BE unit-test function JwtUtilTest::VerifyJwtRS256.
   */
  @Test
  public void testWebserverOAuthAuth() throws Exception {
    String jwksFilename =
        new File(System.getenv("IMPALA_HOME"), "testdata/jwt/jwks_rs256.json").getPath();
    setUp(String.format(
              "--oauth_token_auth=true --oauth_jwt_validate_signature=true "
                  + "--oauth_jwks_file_path=%s --oauth_allow_without_tls=true",
              jwksFilename),
        "");

    // Case 1: Authenticate with valid OAuth Token in HTTP header.
    String oauthToken =
        "eyJhbGciOiJSUzI1NiIsImtpZCI6InB1YmxpYzpjNDI0YjY3Yi1mZTI4LTQ1ZDctYjAxNS1m"
        + "NzlkYTUwYjViMjEiLCJ0eXAiOiJKV1MifQ.eyJpc3MiOiJhdXRoMCIsInVzZXJuYW1lIjoia"
        + "W1wYWxhIn0.OW5H2SClLlsotsCarTHYEbqlbRh43LFwOyo9WubpNTwE7hTuJDsnFoVrvHiWI"
        + "02W69TZNat7DYcC86A_ogLMfNXagHjlMFJaRnvG5Ekag8NRuZNJmHVqfX-qr6x7_8mpOdU55"
        + "4kc200pqbpYLhhuK4Qf7oT7y9mOrtNrUKGDCZ0Q2y_mizlbY6SMg4RWqSz0RQwJbRgXIWSgc"
        + "bZd0GbD_MQQ8x7WRE4nluU-5Fl4N2Wo8T9fNTuxALPiuVeIczO25b5n4fryfKasSgaZfmk0C"
        + "oOJzqbtmQxqiK9QNSJAiH2kaqMwLNgAdgn8fbd-lB1RAEGeyPH8Px8ipqcKsPk0bg";
    attemptConnection("Bearer " + oauthToken, "127.0.0.1");
    verifyAuthMetrics(Range.closed(1L, 1L), zero, "oauth");

    // Case 2: Failed with invalid OAuth Token.
    String invalidOAuthToken =
        "eyJhbGciOiJSUzI1NiIsImtpZCI6InB1YmxpYzpjNDI0YjY3Yi1mZTI4LTQ1ZDctYjAxNS1m"
        + "NzlkYTUwYjViMjEiLCJ0eXAiOiJKV1MifQ.eyJpc3MiOiJhdXRoMCIsInVzZXJuYW1lIjoia"
        + "W1wYWxhIn0.";
    try {
      attemptConnection("Bearer " + invalidOAuthToken, "127.0.0.1");
    } catch (IOException e) {
      assertTrue(e.getMessage().contains("Server returned HTTP response code: 401"));
    }
    verifyAuthMetrics(Range.closed(1L, 1L), Range.closed(1L, 1L), "oauth");

    // Case 3: Failed without "Bearer" token.
    try {
      attemptConnection("Basic VGVzdDFMZGFwOjEyMzQ1", "127.0.0.1");
    } catch (IOException e) {
      assertTrue(e.getMessage().contains("Server returned HTTP response code: 401"));
    }
    // OAUth authentication is not invoked.
    verifyAuthMetrics(Range.closed(1L, 1L), Range.closed(1L, 1L), "oauth");

    // Case 4: Failed without "Authorization" header.
    try {
      attemptConnection(null, "127.0.0.1");
    } catch (IOException e) {
      assertTrue(e.getMessage().contains("Server returned HTTP response code: 401"));
    }
    // OAuth authentication is not invoked.
    verifyAuthMetrics(Range.closed(1L, 1L), Range.closed(1L, 1L), "oauth");
  }

  // Helper method to make a "get" call to the Web Server using the input OAuth auth token
  // and x-forward-for address.
  private void attemptConnection(String auth_token, String xff_address) throws Exception {
    String url = "http://localhost:25000/?json";
    URLConnection connection = new URL(url).openConnection();
    if (auth_token != null) connection.setRequestProperty("Authorization", auth_token);
    if (xff_address != null) {
      connection.setRequestProperty("X-Forwarded-For", xff_address);
    }
    connection.getInputStream();
  }
}
